from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr
from hausdorff import hausdorff_distance

# trj1 = np.array([[0,1],[1,1],[2,1],[2,2],[2,3],[2,4]])
# trj2 = np.array([[0,1],[1,1],[2,1],[2,2],[3,2],[3,3]])
# trj3 = np.array([[0,1],[1,1],[2,1],[2,2],[3,2],[4,2]])
trj1 = np.array([[0,1],[1,1],[2,1],[2,2],])
trj2 = np.array([[0,1],[1,1],[2,1],[2,2],[3,2]])
trj3 = np.array([[0,1],[1,1],[2,1],[2,2],[3,2],[4,2]])

@jax.jit
def euclidean_dist(x1,x2):
    return jnp.sqrt(jnp.sum((x1-x2)**2))
euclidean_dist_vmap = jax.vmap(euclidean_dist, in_axes=(0, None))

@jax.jit
def Frechet_distance(processed_trj1_, processed_trj2_):
    # 将 processed_trj1 和 processed_trj2 各自对齐到它们的起点
    processed_trj1 = processed_trj1_ - processed_trj1_[0] + jnp.array([10,10])
    processed_trj2 = processed_trj2_ - processed_trj2_[0] + jnp.array([10,10])
    # 计算processed_trj1和processed_trj2 中的每个元素之间的欧氏距离，并输出一个长度为max_length的距离列表
    dist_vector = euclidean_dist_vmap(processed_trj1, processed_trj2)
    # 返回 dist_vector 当中的最大值
    return jnp.max(dist_vector)

# 计算 trj1 和 trj2 之间的 Hausdorff 距离
H12 = hausdorff_distance(trj1, trj2, distance='euclidean')
print(f"Hausdorff distance test: {H12}")
# 计算 trj1 和 trj3 之间的 Hausdorff 距离
H13 = hausdorff_distance(trj1, trj3, distance='euclidean')
print(f"Hausdorff distance test: {H13}")

# 计算 trj1 和 trj2 之间的 Frechet 距禋
F12 = Frechet_distance(trj1, trj2)
print(f"Frechet distance test: {F12}")
# 计算 trj1 和 trj3 之间的 Frechet 距离
F13 = Frechet_distance(trj1, trj3)
print(f"Frechet distance test: {F13}")

# 将 H12/H13/F12/F13 绘制成 bar chart, 并写明标注
fig, ax = plt.subplots()
ax.bar([0,1,2,3], [H12, H13, F12, F13])
# 为 H 组 和 F 组设置不同的颜色
ax.get_children()[0].set_color('red')
ax.get_children()[1].set_color('blue')
ax.get_children()[2].set_color('red')
ax.get_children()[3].set_color('blue')
ax.set_xticks([0,1,2,3])
ax.set_xticklabels(['Hausdorff-12', 'Hausdorff-13', 'Frechet-12', 'Frechet-13'])
plt.show()
